import numpy as np
import matplotlib.pyplot as plt

# Funciones Optimizadas (modo vectorial): -----------------------------------------------------------------
def euler(X, t_i, t_f, dt, F):
    t = t_i
    iteraciones = int((t_f - t_i)/dt)
    for i in range(iteraciones):
        X = X + dt*F(X, t)
        t += dt
        for j in range(len(X)):
            plt.figure(j+1)
            plt.title("Método de Euler")
            plt.plot(t, X[j], ".")
            plt.xlabel("t")
            plt.ylabel(f"x{j}")
    plt.show()

def RK2(X, t_i, t_f, dt, F):
    t = t_i
    iteraciones = int((t_f - t_i)/dt)
    for i in range(iteraciones):
        k1 = dt *F(X, t)
        k2 = dt*F(X + k1/2, t + dt/2)
        X = X + k2
        t += dt
        for j in range(len(X)):
            plt.figure(j+1)
            plt.title("Método de Runge-Kutta de 2º orden")
            plt.plot(t, X[j], ".")
            plt.xlabel("t")
            plt.ylabel(f"x{j}")
    plt.show()

def RK4(X, t_i, t_f, dt, F):
    t = t_i
    iteraciones = int((t_f - t_i)/dt)
    for i in range(iteraciones):
        k1 = dt*F(X, t)
        k2 = dt*F(X + k1/2, t + dt/2)
        k3 = dt*F(X + k2/2, t + dt/2)
        k4 = dt*F(X + k3, t + dt)
        X = X + k1/6 + k2/3 + k3/3 + k4/6
        t += dt
        for j in range(len(X)):
            plt.figure(j+1)
            plt.title("Método de Runge-Kutta de 4º orden")
            plt.plot(t, X[j], ".")
            plt.ylabel(f"x{j}")
    plt.show()

"""
Ejemplo X = np.array([x, y, z])      En condiciones iniciales, de cada varfiable

def F(X, t):
    x, y, z = X
    
    dxdt = z
    dydt = -x
    dzdt = y
    
    return np.array([dxdt, dydt, dzdt])

Ejemplo llamar funcion: euler(np.array([(x0, y0)]), t_0, t_f, dt, f(x,y))
"""

# Ejercicio 1 -----------------------------------------------------------------
print("Ejercicio 1 " + "-"*60)

def f(P, t):
    r = 1
    k = 1
    return r*P*(1-(P/k))

condiciones_iniciales, t_i, t_f, dt = np.array([10]), 0, 10, 0.01

euler(condiciones_iniciales, t_i, t_f, dt, f)
RK2(condiciones_iniciales, t_i, t_f, dt, f)
RK4(condiciones_iniciales, t_i, t_f, dt, f)

# Ejercicio 2 -----------------------------------------------------------------
print("Ejercicio 2 " + "-"*60)

def oscilador(X, t): 
    x, y = X
    
    w_0 = 1
    dxdt = y
    dydt = -w_0**2 *x
    return np.array([dxdt, dydt])

condiciones_iniciales, t_i, t_f, dt = np.array([1, 0]), 0, 10, 0.01

euler(condiciones_iniciales, t_i, t_f, dt, oscilador)
RK2(condiciones_iniciales, t_i, t_f, dt, oscilador)
RK4(condiciones_iniciales, t_i, t_f, dt, oscilador)

# Ejercicio 3 -----------------------------------------------------------------
print("Ejercicio 3 " + "-"*60)

def oscilador_amortiguado(X, t):
    x, y = X
    
    dxdt = y
    dydt = F*np.cos(w*t)-w_0**2 *x - b*y
    return np.array([dxdt, dydt])

condiciones_iniciales, t_i, t_f, dt = np.array([3, 2]), 0, 10, 0.01

w_0, w, b, F = 1, 1, 0.2, 1
euler(condiciones_iniciales, t_i, t_f, dt, oscilador_amortiguado)
RK2(condiciones_iniciales, t_i, t_f, dt, oscilador_amortiguado)
RK4(condiciones_iniciales, t_i, t_f, dt, oscilador_amortiguado)

w_0, w, b, F = 1, 0.8, 2, 0.5
euler(condiciones_iniciales, t_i, t_f, dt, oscilador_amortiguado)
RK2(condiciones_iniciales, t_i, t_f, dt, oscilador_amortiguado)
RK4(condiciones_iniciales, t_i, t_f, dt, oscilador_amortiguado)

# Ejercicio 4 -----------------------------------------------------------------
print("Ejercicio 4 " + "-"*60)

def sistema_ED(X, t):
    x, y, z = X
    sigma, r, b = 3, 26.5, 1
    
    dxdt = sigma*(y - x)
    dydt = r*x - y - x*z
    dzdt = x*y - b * z
    
    return np.array([dxdt, dydt, dzdt])

condiciones_iniciales, t_i, t_f, dt = np.array([0, 1, 0]), 0, 10, 0.01
euler(condiciones_iniciales, t_i, t_f, dt, sistema_ED)
RK2(condiciones_iniciales, t_i, t_f, dt, sistema_ED)
RK4(condiciones_iniciales, t_i, t_f, dt, sistema_ED)

# Ejercicio 5 -----------------------------------------------------------------
print("Ejercicio 5 " + "-"*60)

condiciones_iniciales, t_i, t_f, dt = np.array([10]), 0, 10, 0.01

# Crear método de paso variable
"""
def metodo_paso_variable(X, t_i, t_f, dt, F):
    t = t_i
    
    dt_peque, dt_mediana, dt_grande = dt/20, dt, 20*dt 
    limitador_pendiente = 1     # Ponemos que entre 0 y 25% es pendiente pequeña, entre 10% y 100% pendiente mediana y >100% es pendiente grade

    while t < t_f:
        X1 = X +dt*F(X, t)
        
        pendiente = abs(F(X, t))
        #dt = dt_mediana/pendiente
        
        X = X+dt*F(X, t)
        t += dt
        for j in range(len(X)):
            plt.title("Método de Paso Variable")
            plt.figure(j+1)
            plt.plot(t, X[j], ".")
            plt.xlabel("t")
            plt.ylabel(f"x{j}")
    plt.show()
"""
def metodo_paso_variable(X, t_i, t_f, dt, F, tol=1e-4):
    t = t_i
    dt_min, dt_max = dt/1000, dt*100

    while t < t_f:
        if t + dt > t_f:
            dt = t_f - t

        X_euler = X + dt * F(X, t)
        k1 = dt * F(X, t)
        k2 = dt * F(X + k1/2, t + dt/2)
        X_rk2 = X + k2

        error = np.linalg.norm(X_rk2 - X_euler)

        if error < tol or dt <= dt_min:
            X = X_rk2          # acepta el paso
            t += dt
            for j in range(len(X)):
                plt.title("Método de Paso Variable")
                plt.figure(j+1)
                plt.plot(t, X[j], ".")
                plt.xlabel("t")
                plt.ylabel(f"x{j}")

        # Ajusta dt para el siguiente paso
        if error > 0:
            dt = np.clip(0.9 * dt * (tol/error)**0.5, dt_min, dt_max)
    plt.show()

condiciones_iniciales, t_i, t_f, dt = np.array([10]), 0, 10, 0.01

metodo_paso_variable(condiciones_iniciales, t_i, t_f, dt, f)

# Ejercicio 6 -----------------------------------------------------------------
print("Ejercicio 6 " + "-"*60)

# APARTADO a)
def oscilador_fuerza_lineal(X, t):
    x, y = X

    w_0, b, F = 5, 0.2, 10

    dxdt = y
    dydt = F*t - b*y - w_0**2 *x
    
    return np.array([dxdt, dydt])

condiciones_iniciales, t_i, t_f, dt = np.array([1, 1]), 0, 10, 0.01

euler(condiciones_iniciales, t_i, t_f, dt, oscilador_fuerza_lineal)

# APARTADO b)

def euler_datos(X0, t_i, t_f, dt, F):
    t_vals = np.arange(t_i, t_f, dt)
    X = np.array(X0)

    X_vals = [X.copy()]

    for t in t_vals[:-1]:
        X = X + dt * F(X, t)
        X_vals.append(X.copy())

    return t_vals, np.array(X_vals)

t, X = euler_datos(condiciones_iniciales, 0, 10, 0.01, oscilador_fuerza_lineal)

x = X[:,0]
v = X[:,1]
# Aqui se almacenan las listas de t_max y x_max
t_max = []
x_max = []

for i in range(len(v)-1):
    if v[i] > 0 and v[i+1] < 0:
        t_max.append(t[i])
        x_max.append(x[i])
print(f"x_max = {x_max}\nt_max = {t_max}")

plt.plot(t, x, label="x(t)")
plt.xlabel("t")
plt.ylabel("x")
plt.plot(t_max, x_max, "ro", label="máximos")
plt.legend()
plt.show()

# APARTADO c)
coef = np.polyfit(t_max, x_max, 1)

a, b = coef
print("\nPendiente a =", a)
print("Ordenada b =", b)

t_fit = np.linspace(min(t_max), max(t_max), 100)
x_fit = a*t_fit + b

plt.xlabel("t")
plt.ylabel("x")
plt.plot(t_max, x_max, "ro", label="máximos")
plt.plot(t_fit, x_fit, label=f"ajuste lineal: y={a:.2f}·t+{b:.2f}")
plt.legend()
plt.show()